import torch
import torch.nn as nn

class L1L2Loss(nn.Module):
    def __init__(self, args):
        super(L1L2Loss, self).__init__()

        self.max_depth = args.max_depth
        self.min_depth = args.min_depth
        self.t_valid = self.min_depth

    # L1L2组合损失
    def l1l2comb(self, pred, gt):
        mask = (gt > self.t_valid).type_as(pred).detach() # 二值监督掩码
        num_valid = torch.sum(mask, dim=[1, 2, 3]) # 有效值数量

        d1 = torch.sum(torch.abs(pred - gt) * mask, dim=[1, 2, 3]) # L1损失
        d2 = torch.sum(torch.pow(pred - gt, 2) * mask, dim=[1, 2, 3]) # L2损失

        return (d1+d2) * 0.1 / (num_valid + 1e-8)

    def forward(self, output, gt, temper=0.1):
        """
        pred , gt = torch.Size([B, 1, H, W])
        """

        pred = output['pred'] # 深度图[b,1,h,w]
        init = output['pred_init'] # 相对深度[b,1,h,w]

        # 深度图截断
        gt = torch.clamp(gt, min=0, max=self.max_depth)
        pred = torch.clamp(pred, min=self.min_depth, max=self.max_depth)
        init = torch.clamp(init, min=self.min_depth, max=self.max_depth)

        # 伪标签
        gt_p = torch.where(torch.rand_like(init) > temper, torch.tensor(self.min_depth).to(init), init) # 随机取点做为伪标签，概率0.1
        gt_p = torch.where(init == self.max_depth, torch.tensor(self.max_depth).to(gt_p), gt_p) # 添加天空区域
        gt_p = gt_p * (gt < self.min_depth) # 剔除真实标签

        l1l2_loss = self.l1l2comb(pred, gt) # 稀疏点监督
        l1l2_loss_p = self.l1l2comb(pred, gt_p) # 伪标签监督

        loss = l1l2_loss + 0.1 * l1l2_loss_p # 组合损失

        return loss.mean()
